import sys
import cv2
import numpy as np
from PyQt5.QtWidgets import (
    QApplication, QMainWindow, QFileDialog, QLabel, QVBoxLayout, QPushButton, QWidget
)
from PyQt5.QtGui import QPixmap, QImage, QPainter
from PyQt5.QtCore import Qt, QRect

class DocCorrectionApp(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("文档矫正程序")
        self.setGeometry(100, 100, 800, 600)
        self.image_label = QLabel(self)
        self.image_label.setScaledContents(True)
        self.image_label.setMouseTracking(True)

        self.init_ui()
        self.image = None
        self.rect_start = None
        self.rect_end = None
        self.rect_region = None

    def init_ui(self):
        # 主布局
        layout = QVBoxLayout()

        # 按钮
        self.load_button = QPushButton("加载图片")
        self.correct_button = QPushButton("矫正文档")
        self.correct_button.setEnabled(False)

        # 信号连接
        self.load_button.clicked.connect(self.load_image)
        self.correct_button.clicked.connect(self.correct_document)

        # 添加到布局
        layout.addWidget(self.load_button)
        layout.addWidget(self.correct_button)
        layout.addWidget(self.image_label)

        # 设置主窗口布局
        container = QWidget()
        container.setLayout(layout)
        self.setCentralWidget(container)

    def load_image(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "选择图片文件", "", "Images (*.png *.jpg *.jpeg *.bmp)")
        if file_path:
            self.image = cv2.imread(file_path)
            self.display_image(self.image)
            self.correct_button.setEnabled(True)

    def display_image(self, img):
        """将 OpenCV 图像显示在 QLabel 上"""
        rgb_image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        h, w, ch = rgb_image.shape
        bytes_per_line = ch * w
        q_image = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
        self.image_label.setPixmap(QPixmap.fromImage(q_image))
        self.image_label.setFixedSize(w, h)

    def correct_document(self):
        """对选定区域进行文档矫正"""
        if self.rect_region is None or self.image is None:
            return

        x1, y1, x2, y2 = self.rect_region.getCoords()
        cropped = self.image[y1:y2, x1:x2]

        # 透视变换：假设目标为一个 A4 长宽比的矩形
        h, w = cropped.shape[:2]
        src_pts = np.array([[0, 0], [w - 1, 0], [w - 1, h - 1], [0, h - 1]], dtype=np.float32)
        dst_pts = np.array([[0, 0], [210, 0], [210, 297], [0, 297]], dtype=np.float32)

        matrix = cv2.getPerspectiveTransform(src_pts, dst_pts)
        corrected = cv2.warpPerspective(cropped, matrix, (210, 297))

        self.display_image(corrected)

    def mousePressEvent(self, event):
        """鼠标按下事件：记录起始点"""
        if self.image is None:
            return
        if event.button() == Qt.LeftButton:
            self.rect_start = event.pos()

    def mouseReleaseEvent(self, event):
        """鼠标释放事件：记录结束点"""
        if self.image is None:
            return
        if event.button() == Qt.LeftButton and self.rect_start is not None:
            self.rect_end = event.pos()
            self.rect_region = QRect(self.rect_start, self.rect_end)
            self.update()

    def paintEvent(self, event):
        """绘制选定矩形"""
        super().paintEvent(event)
        if self.rect_region:
            painter = QPainter(self)
            painter.setPen(Qt.red)
            painter.drawRect(self.rect_region)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    main_win = DocCorrectionApp()
    main_win.show()
    sys.exit(app.exec_())
